package com.ming.util;

import lombok.extern.slf4j.Slf4j;
import org.slf4j.MDC;

import java.util.UUID;

/**
 * TraceId工具类
 * 用于生成和管理请求追踪ID
 *
 * @author ming
 * @version 1.0
 */
@Slf4j
public class TraceIdUtils {

    private static final String TRACE_ID_KEY = "traceId";

    /**
     * 生成TraceId
     */
    public static String generateTraceId() {
        return UUID.randomUUID().toString().replace("-", "");
    }

    /**
     * 设置TraceId到MDC
     */
    public static void setTraceId(String traceId) {
        if (traceId != null) {
            MDC.put(TRACE_ID_KEY, traceId);
        }
    }

    /**
     * 生成并设置TraceId
     */
    public static String generateAndSetTraceId() {
        String traceId = generateTraceId();
        setTraceId(traceId);
        return traceId;
    }

    /**
     * 获取当前TraceId
     */
    public static String getTraceId() {
        return MDC.get(TRACE_ID_KEY);
    }

    /**
     * 清除TraceId
     */
    public static void clearTraceId() {
        MDC.remove(TRACE_ID_KEY);
    }

    /**
     * 清除MDC
     */
    public static void clear() {
        MDC.clear();
    }

    /**
     * 如果TraceId不存在，则生成并设置
     */
    public static String ensureTraceId() {
        String traceId = getTraceId();
        if (traceId == null) {
            traceId = generateAndSetTraceId();
        }
        return traceId;
    }

    /**
     * 执行任务并自动管理TraceId
     */
    public static <T> T runWithTraceId(String traceId, Runnable task) {
        String originalTraceId = getTraceId();
        try {
            setTraceId(traceId);
            task.run();
            return null;
        } finally {
            if (originalTraceId != null) {
                setTraceId(originalTraceId);
            } else {
                clearTraceId();
            }
        }
    }

    /**
     * 执行任务并自动管理TraceId（有返回值）
     */
    public static <T> T runWithTraceId(String traceId, java.util.function.Supplier<T> supplier) {
        String originalTraceId = getTraceId();
        try {
            setTraceId(traceId);
            return supplier.get();
        } finally {
            if (originalTraceId != null) {
                setTraceId(originalTraceId);
            } else {
                clearTraceId();
            }
        }
    }
}